include(cc_library)
include(cc_test)

cc_library(
  NAME 
    sampler
  HDRS
    sampling_params.h
    logits_utils.h
    rejection_sampler.h
    sampler.h
  SRCS
    sampling_params.cpp
    logits_utils.cpp
    rejection_sampler.cpp
    sampler.cpp
  DEPS
    glog::glog
    torch
    $<$<BOOL:${USE_NPU}>:xllm_ops>
)

cc_test(
  NAME
    sampler_test
  SRCS
    rejection_sampler_test.cpp
    rejection_sampler.cpp
    sampling_params_test.cpp
  DEPS
    absl::strings
    GTest::gtest_main
    :flags
    :sampler
    glog::glog
)
target_link_libraries(sampler_test PRIVATE brpc OpenSSL::SSL OpenSSL::Crypto leveldb::leveldb ZLIB::ZLIB protobuf::libprotobuf)
target_link_libraries(sampler_test
                      PUBLIC
                      Python::Python
                      $<$<BOOL:${USE_NPU}>:ascendcl>
                      $<$<BOOL:${USE_NPU}>:hccl>
                      $<$<BOOL:${USE_NPU}>:c_sec>
                      $<$<BOOL:${USE_NPU}>:nnopbase>)
add_dependencies(sampler_test brpc-static)